package org.example.rag;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author fish_temp_author
 * @since fish_temp_since
 */
public class SimpleRAG {
    public static void main(String[] args) {

        OpenAiEmbeddingModel openAiEmbeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey("demo")
                .modelName("text-embedding-3-small")
                .build();

        RedisEmbeddingStore redisEmbeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .indexName("rag")
                // key前缀
                .prefix("rag:vec:")
                // 维度 当前向量模型返回的向量是1536维度
                .dimension(1536)
                .build();

        OpenAiChatModel openAiChatModel = OpenAiChatModel.builder()
                .apiKey("demo")
                .modelName("gpt-4o-mini")
                .build();

        // 1. 对知识通过向量模型进行向量化，存储到向量库中
        TextSegment textSegment = TextSegment.textSegment("客服的名字叫谢阿花");
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment).content(), textSegment);
        TextSegment textSegment2 = TextSegment.textSegment("客服的电话是：0791-63213328");
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment2).content(), textSegment2);
        TextSegment textSegment3 = TextSegment.textSegment("客服的服务时间是9:00 - 20:00");
        redisEmbeddingStore.add(openAiEmbeddingModel.embed(textSegment3).content(), textSegment3);

        // 2. 对用户的问题进行向量化，并查询数据库中相似度最高的向量以及对应的知识内容
        String q = "客服的电话和名字以及工作时间";
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(openAiEmbeddingModel.embed(q).content())
                .minScore(0.6d)
                .maxResults(3)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = redisEmbeddingStore.search(searchRequest);

        List<EmbeddingMatch<TextSegment>> matches = searchResult.matches();
        if (matches.isEmpty()) {
            System.out.println("没有找到答案");
        } else {
            String acknowledge = matches.stream()
                    .map(x -> x.embedded().text())
                    .collect(Collectors.joining("\n"));

            // 3. 将用户的问题与知识内容进行拼接，作为模型的输入，进行模型预测，得到答案

            PromptTemplate template = PromptTemplate.from("""
                    用户的问题：{{question}}
                    请严格根据上下文<acknowledge>标签中的内容来回答用户的问题，如果上下文为空，请回答“我还在学习中呢～”。
                    <acknowledge>{{acknowledge}}</acknowledge>
                    """);
            // 使用 Map 来传递参数
            Map<String, Object> variables = new HashMap<>();
            variables.put("question", q);
            variables.put("acknowledge", acknowledge);

            Prompt prompt = template.apply(variables);

            System.out.println("prompt:" + prompt.toString());
            System.out.println(openAiChatModel.chat(prompt.toUserMessage())
                    .aiMessage()
                    .text());
        }
    }
}
